# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# croppping utilities
# --------------------------------------------------------
import sys
import PIL.Image
import os

from PIL import Image
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2  # noqa
import numpy as np  # noqa

try:
    lanczos = PIL.Image.Resampling.LANCZOS
    bicubic = PIL.Image.Resampling.BICUBIC
except AttributeError:
    lanczos = PIL.Image.LANCZOS
    bicubic = PIL.Image.BICUBIC

# from resize_hypersim
def crop_resize_if_necessary(image, depthmap, intrinsics, resolution, info=None, crop=True):
    """
    Crop and resize the given image, depthmap, and intrinsics.
    
    The procedure is as follows:
      1. Ensure the image is a PIL.Image.Image.
      2. Determine a centered crop window on the principal point.
      3. Crop the image and depthmap.
      4. Downscale the result using Lanczos interpolation so that the image size matches resolution.
      5. Compute new intrinsics from the cropping.
      6. Perform final cropping (if necessary) with bilinear interpolation.
    
    Args:
        image: an RGB image (PIL Image or numpy array)
        depthmap: a depth map (numpy array)
        intrinsics: a 3x3 numpy array containing camera intrinsics.
        resolution: a tuple (width, height) specifying target resolution.
        info: optional info for error messages.
    
    Returns:
        A tuple of (image_out, depthmap_out, intrinsics_out)
    """
    # Convert image to PIL image if needed.
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    W, H = image.size
    # print(W, H)
    # Use the principal point from intrinsics (assumed to be in position [0,2] and [1,2])
    if crop:
        cx, cy = intrinsics[:2, 2] + 0.5
        min_margin_x = min(cx, W - cx)
        min_margin_y = min(cy, H - cy)
        if min_margin_x <= W / 5 or min_margin_y <= H / 5:
            raise ValueError(f"Bad principal point in view {info}")
        # Define crop bounding box centered on (cx, cy)
        # l, t = cx - min_margin_x, cy - min_margin_y
        # r, b = cx + min_margin_x, cy + min_margin_y
        crop_width = int(min_margin_x) * 2
        crop_height = int(min_margin_y) * 2
        l = int(round(cx - crop_width / 2.0))
        t = int(round(cy - crop_height / 2.0))
        r = l + crop_width
        b = t + crop_height
        crop_bbox = (l, t, r, b)
        # print(crop_bbox, W, H, crop_width, crop_height, cx, cy, min_margin_x, min_margin_y)
        image, depthmap, intrinsics = crop_image_depthmap(image, depthmap, intrinsics, crop_bbox)
        # print('depthmap', depthmap.shape, image.size, resolution)
    # Downscale using high-quality Lanczos interpolation so that image.size == resolution.
    if isinstance(resolution, int):
        target_w = resolution
        target_h = resolution
    else:
        target_w = resolution[0]
        target_h = resolution[1]
    # print(target_w, target_h)
    image, depthmap, intrinsics = rescale_image_depthmap(image, depthmap, intrinsics, (target_w, target_h))
    W, H = image.size
    resolution = image.size
    # print(image.size, depthmap.shape)
    # sys.exit(0)
    # Compute new intrinsics from final cropping.
    intrinsics2 = camera_matrix_of_crop(intrinsics, image.size, resolution, offset_factor=0.5)
    crop_bbox = bbox_from_intrinsics_in_out(intrinsics, intrinsics2, resolution)
    image, depthmap, intrinsics2 = crop_image_depthmap(image, depthmap, intrinsics, crop_bbox)
    # print(image.size, depthmap.shape, intrinsics2.shape)
    
    return image, depthmap, intrinsics2

def colmap_to_opencv_intrinsics(K):
    """
    Modify camera intrinsics to follow a different convention.
    Coordinates of the center of the top-left pixels are by default:
    - (0.5, 0.5) in Colmap
    - (0,0) in OpenCV
    """
    K = K.copy()
    K[0, 2] -= 0.5
    K[1, 2] -= 0.5
    return K


def opencv_to_colmap_intrinsics(K):
    """
    Modify camera intrinsics to follow a different convention.
    Coordinates of the center of the top-left pixels are by default:
    - (0.5, 0.5) in Colmap
    - (0,0) in OpenCV
    """
    K = K.copy()
    K[0, 2] += 0.5
    K[1, 2] += 0.5
    return K


class ImageList:
    """Convenience class to aply the same operation to a whole set of images."""

    def __init__(self, images):
        if not isinstance(images, (tuple, list, set)):
            images = [images]
        self.images = []
        for image in images:
            if not isinstance(image, PIL.Image.Image):
                image = PIL.Image.fromarray(image)
            self.images.append(image)

    def __len__(self):
        return len(self.images)

    def to_pil(self):
        return tuple(self.images) if len(self.images) > 1 else self.images[0]

    @property
    def size(self):
        sizes = [im.size for im in self.images]
        assert all(sizes[0] == s for s in sizes)
        return sizes[0]

    def resize(self, *args, **kwargs):
        return ImageList(self._dispatch("resize", *args, **kwargs))

    def crop(self, *args, **kwargs):
        return ImageList(self._dispatch("crop", *args, **kwargs))

    def _dispatch(self, func, *args, **kwargs):
        return [getattr(im, func)(*args, **kwargs) for im in self.images]


def rescale_image_depthmap(
    image, depthmap, camera_intrinsics, output_resolution, force=True
):
    """Jointly rescale a (image, depthmap)
    so that (out_width, out_height) >= output_res
    """
    image = ImageList(image)
    input_resolution = np.array(image.size)  # (W,H)
    output_resolution = np.array(output_resolution)
    # print(input_resolution, output_resolution, depthmap.shape, image.size[::-1])
    # if depthmap is not None:
        # can also use this with masks instead of depthmaps
        # assert tuple(depthmap.shape[:2]) == image.size[::-1]
    # define output resolution
    assert output_resolution.shape == (2,)
    scale_final = max(output_resolution / image.size) + 1e-8
    if scale_final >= 1 and not force:  # image is already smaller than what is asked
        return (image.to_pil(), depthmap, camera_intrinsics)
    output_resolution = np.floor(input_resolution * scale_final).astype(int)

    # first rescale the image so that it contains the crop
    # print(output_resolution, scale_final)
    image = image.resize(
        tuple(output_resolution), resample=lanczos if scale_final < 1 else bicubic
    )
    # print(image.size)
    if depthmap is not None:
        depthmap = cv2.resize(
            depthmap,
            output_resolution,
            fx=scale_final,
            fy=scale_final,
            interpolation=cv2.INTER_NEAREST,
        )
    # print(depthmap.shape)
    # no offset here; simple rescaling
    camera_intrinsics = camera_matrix_of_crop(
        camera_intrinsics, input_resolution, output_resolution, scaling=scale_final
    )

    return image.to_pil(), depthmap, camera_intrinsics


def camera_matrix_of_crop(
    input_camera_matrix,
    input_resolution,
    output_resolution,
    scaling=1,
    offset_factor=0.5,
    offset=None,
):
    # Margins to offset the origin
    margins = np.asarray(input_resolution) * scaling - output_resolution
    assert np.all(margins >= 0.0)
    if offset is None:
        offset = offset_factor * margins

    # Generate new camera parameters
    output_camera_matrix_colmap = opencv_to_colmap_intrinsics(input_camera_matrix)
    output_camera_matrix_colmap[:2, :] *= scaling
    output_camera_matrix_colmap[:2, 2] -= offset
    output_camera_matrix = colmap_to_opencv_intrinsics(output_camera_matrix_colmap)

    return output_camera_matrix


def crop_image_depthmap(image, depthmap, camera_intrinsics, crop_bbox):
    """
    Return a crop of the input view.
    """
    image = ImageList(image)
    l, t, r, b = crop_bbox

    image = image.crop((l, t, r, b))
    depthmap = depthmap[t:b, l:r]

    camera_intrinsics = camera_intrinsics.copy()
    camera_intrinsics[0, 2] -= l
    camera_intrinsics[1, 2] -= t

    return image.to_pil(), depthmap, camera_intrinsics


def bbox_from_intrinsics_in_out(
    input_camera_matrix, output_camera_matrix, output_resolution
):
    out_width, out_height = output_resolution
    l, t = np.int32(np.round(input_camera_matrix[:2, 2] - output_camera_matrix[:2, 2]))
    crop_bbox = (l, t, l + out_width, t + out_height)
    return crop_bbox
